#include "Crypter.h"
#include "assert.h"
#pragma comment(lib, "advapi32.lib")

DataBlock::DataBlock():data(0),data_len(0)
{
}
DataBlock::~DataBlock()
{
    Free();
}

DataBlock::DataBlock(unsigned long size):data(0),data_len(0)
{
    Create(size);
}

void DataBlock::Create(unsigned long size)
{
    if(data)
    {
        if(data_len!=size)
        {
            data_len=size;
            data=(unsigned char*)realloc(data,data_len);

          }
    }
    else
    {
        data_len=size;
        data =(unsigned char*)malloc(data_len);
    }

}
void DataBlock::Add(const unsigned char* _data,const unsigned long size)
{
    unsigned long len=data_len;
    data_len+=size;
    data=(unsigned char*)realloc(data,data_len);
    memcpy(data+len,_data,size);
}

void DataBlock::CreateFrom(const unsigned char* _data,const unsigned long size)
{
    Create(size);
    memcpy(data,_data,size);
}
void DataBlock::CreateFromBlock(DataBlock& block)
{
    quint32 size=block.GetLength();
    Create(size);
    memcpy(data,block.Get(),size);
}

void DataBlock::CreateFromStr(const char* _data,bool zero)
{
    quint32 size=strlen(_data);
    Create(size+zero);
    memcpy(data,_data,size);
    if(zero)
        data[size]=0;
}



void DataBlock::Free()
{
    if(data)
    {
        free(data);
        data=0;
        data_len=0;
    }
}
unsigned long DataBlock::GetLength()
{
    return data_len;
}

unsigned char* DataBlock::Get(unsigned long* size)
{
    if(size)
        *size=data_len;
    return data;
}

Crypter::Crypter(int size):isRSAInit(false),isAESInit(false),mDataBuffer(0),mIVCounter(0)
{
     mDataBuffer=new unsigned char[size];

        Secure(register_prng(&fortuna_desc));
        Secure(register_hash(&sha256_desc));
        Secure(register_cipher(&aes_desc));

        hash_idx = find_hash("sha256");
        prng_idx = find_prng("fortuna");
        cipher_idx =find_cipher("aes");
        ltc_mp = ltm_desc;
        unsigned long t=QTime(0,0,0).secsTo(QTime::currentTime());

         Secure(fortuna_start(&mPrngState));
         Secure(fortuna_add_entropy((unsigned char*)&t,sizeof(t),&mPrngState));
         Secure(fortuna_ready(&mPrngState));





}
bool Crypter::Secure()
{
    return true;
}

bool Crypter::Secure(int error)
{
    if(error==CRYPT_OK)
        return true;
    else
       {
            std::string msg=error_to_string(error);
            return false;
        }

}

Crypter::~Crypter()
{
    delete[] mDataBuffer;
    mDataBuffer=0;

    if(isRSAInit)
    {
        rsa_free(&mCipherKey);
        rsa_free(&mSignKey);
        isRSAInit=false;
    }
    if(isAESInit)
    {
        unsigned char tag[256];unsigned long taglen=256;
        //Secure(gcm_done(&gcm,tag,&taglen));
    }
}
void Crypter::SetIVCounter(int counter)
{
    if(isAESInit)
        mIVCounter=counter*10;

}
void Crypter::encode(TCipherType type, DataBlock& source,DataBlock& dest)
{
    switch(type)
    {
        case RSA:
        {
            if(isRSAInit)
            {
                unsigned long dest_size=MAX_BUFFER_SIZE;
                if(Secure(rsa_encrypt_key(source.Get(),source.GetLength(),mDataBuffer,&dest_size,
                                0,0,&mPrngState,prng_idx,hash_idx,&mCipherKey)))
                {
                    dest.CreateFrom(mDataBuffer,dest_size);
                    return ;
                }
                return ;

            }
            else
                return ;
        }
        break;
        case AES:
        {
            if(isAESInit)
            {
                unsigned char tag[BIT128];
                unsigned long tag_len=BIT128;
                Secure(gcm_reset(&gcm));
                Secure(gcm_init(&gcm,cipher_idx,mAESKey,16));
                Secure(gcm_add_iv(&gcm,mAES_IV,16));
                Secure(gcm_add_iv(&gcm,(unsigned char*)&mIVCounter,sizeof(mIVCounter)));
                Secure(gcm_add_aad(&gcm,mAES_AUTH,16));

                int error=Secure(gcm_process(&gcm,source.Get(),source.GetLength(),mDataBuffer,GCM_ENCRYPT));
                dest.CreateFrom(mDataBuffer,source.GetLength());
                gcm_done(&gcm,tag,&tag_len);
                return ;

            }
            else
                return ;
        }
        break;
        case Base64:
            unsigned long len=MAX_BUFFER_SIZE;
            Secure(base64_encode(source.Get(),source.GetLength(),mDataBuffer,&len));
            dest.CreateFrom(mDataBuffer,len+1);
            dest.Get()[len]=0;
        break;

    };
    return ;
}

void Crypter::decode(TCipherType type, DataBlock& source,DataBlock& dest)
{
    switch(type)
    {
        case RSA:
        {
            if(isRSAInit)
            {
                unsigned long dest_size=MAX_BUFFER_SIZE;
                int stat;
                if(Secure(rsa_decrypt_key(source.Get(),source.GetLength(),mDataBuffer,&dest_size,
                                0,0,hash_idx,&stat,&mCipherKey)))
                {
                    dest.CreateFrom(mDataBuffer,dest_size);
                    return ;
                }
            }
            else
                return;
        }
        break;
        case AES:
        {
            if(isAESInit)
            {
                unsigned char tag[BIT128];
                unsigned long tag_len=BIT128;
                Secure(gcm_reset(&gcm));
                Secure(gcm_init(&gcm,cipher_idx,mAESKey,16));
                Secure(gcm_add_iv(&gcm,mAES_IV,16));
                Secure(gcm_add_iv(&gcm,(unsigned char*)&mIVCounter,sizeof(mIVCounter)));
                Secure(gcm_add_aad(&gcm,mAES_AUTH,16));

                int error=Secure(gcm_process(&gcm,mDataBuffer,source.GetLength(),source.Get(),GCM_DECRYPT));
                dest.CreateFrom(mDataBuffer,source.GetLength());
                gcm_done(&gcm,tag,&tag_len);
                 dest.CreateFrom(mDataBuffer,source.GetLength());
                 return;
            }
            else
                return ;
        }
        break;
    case Base64:

        unsigned long len=MAX_BUFFER_SIZE;
        int size=source.GetLength();
        size=source.Get()[size-1]==0?size-1:size;
        Secure(base64_decode(source.Get(),size,mDataBuffer,&len));
        dest.CreateFrom(mDataBuffer,len);

    break;

    };
    return ;
}
void Crypter::calcSign(DataBlock& data, DataBlock& sign)
{
    hash_state md;
    unsigned char bf[BIT256];
    Secure(sha256_init(&md));
    Secure(sha256_process(&md,data.Get(),data.GetLength()));
    Secure(sha256_done(&md,bf));
    int k=fortuna_test();
    unsigned long len=MAX_BUFFER_SIZE;
    Secure(rsa_sign_hash(bf,BIT256,mDataBuffer,&len,&mPrngState,prng_idx,hash_idx,BIT128,&mSignKey));
    sign.CreateFrom(mDataBuffer,len);

}


void Crypter::verifySign(DataBlock& sign, DataBlock& sample,int &correct)
{
    hash_state md;
    unsigned char  bf[BIT256];//256-bit
    Secure(sha256_init(&md));
    Secure(sha256_process(&md,sample.Get(),sample.GetLength()));
    Secure(sha256_done(&md,bf));

    unsigned long len=MAX_BUFFER_SIZE;
    Secure(rsa_verify_hash(sign.Get(),sign.GetLength(),bf,BIT256,hash_idx,16,&correct,&mSignKey));
    return;
}
#include "Satory.h"
void Crypter::GenAESKeys(DataBlock& pass,DataBlock& iv,DataBlock& auth)
{
    if(isAESInit)
    {
        unsigned char tag[256];unsigned long taglen=256;
        Secure(gcm_reset(&gcm));
    }

    string name=Satory::GetInstance()->GetName();
    fortuna_read(mAESKey,16,&mPrngState);
    fortuna_read(mAES_IV,16,&mPrngState);
    fortuna_read(mAES_AUTH,16,&mPrngState);


    isAESInit=true;

    pass.CreateFrom(mAESKey,16);
    iv.CreateFrom(mAES_IV,16);
    auth.CreateFrom(mAES_AUTH,16);
}
void Crypter::AcceptAESKeys(DataBlock& pass,DataBlock& iv,DataBlock& auth)
{
    if(isAESInit)
    {
        unsigned char tag[256];unsigned long taglen=256;
        Secure(gcm_reset(&gcm));
    }
    unsigned long len=MAX_BUFFER_SIZE;
    memcpy(mAESKey,pass.Get(),16);
    memcpy(mAES_IV,iv.Get(),16);
    memcpy(mAES_AUTH,auth.Get(),16);

    isAESInit=true;
}
void Crypter::ResetRSA()
{
    if(isRSAInit)
    {
        rsa_free(&mCipherKey);
        rsa_free(&mSignKey);
        isRSAInit=false;
    }
}
/*void Crypter::GetPublicKeys(DataBlock& cipher,DataBlock& sign   )
{

    unsigned long len=MAX_BUFFER_SIZE;
    Secure(rsa_export(mDataBuffer,&len,PK_PUBLIC,&mCipherKey));
    cipher.CreateFrom(mDataBuffer,len);
    len=MAX_BUFFER_SIZE;
    Secure(rsa_export(mDataBuffer,&len,PK_PUBLIC,&mSignKey));
    sign.CreateFrom(mDataBuffer,len);


}*/
void Crypter::GenRSAKeys(DataBlock& cipher,DataBlock& sign,
                         DataBlock& private_cipher,DataBlock& private_sign)
{
    if(isRSAInit)
    {
        rsa_free(&mCipherKey);
        rsa_free(&mSignKey);
        isRSAInit=false;
    }
    Secure(rsa_make_key(&mPrngState,  prng_idx,  4096/8, 65537, &mCipherKey));
    Secure(rsa_make_key(&mPrngState,  prng_idx,  4096/8, 65537, &mSignKey));
    isRSAInit=true;
    unsigned long len=MAX_BUFFER_SIZE;
    Secure(rsa_export(mDataBuffer,&len,PK_PUBLIC,&mCipherKey));
    cipher.CreateFrom(mDataBuffer,len);
    len=MAX_BUFFER_SIZE;
    Secure(rsa_export(mDataBuffer,&len,PK_PUBLIC,&mSignKey));
    sign.CreateFrom(mDataBuffer,len);
    len=MAX_BUFFER_SIZE;
    Secure(rsa_export(mDataBuffer,&len,PK_PRIVATE,&mCipherKey));
    private_cipher.CreateFrom(mDataBuffer,len);
    len=MAX_BUFFER_SIZE;
    Secure(rsa_export(mDataBuffer,&len,PK_PRIVATE,&mSignKey));
    private_sign.CreateFrom(mDataBuffer,len);

}

void Crypter::AcceptRSAKeys(DataBlock& cipher,DataBlock& sign,bool full)
{
    if(isRSAInit)
    {
        rsa_free(&mCipherKey);
        rsa_free(&mSignKey);
        isRSAInit=false;
    }
    unsigned long len=MAX_BUFFER_SIZE;unsigned char* data=0;

    data=cipher.Get(&len);
    Secure(rsa_import(data,len,&mCipherKey));
    len=MAX_BUFFER_SIZE;
    data=sign.Get(&len);
    Secure(rsa_import(data,len,&mSignKey));

    isRSAInit=true;
}
